#include "gl_memory_pool.h"
#include "gl_define.h"
#include "glcs_glsl_code.h"

namespace PanNN {
namespace OpenGL {
GLMemeoryPool::GLMemeoryPool() {}

GLMemeoryPool::~GLMemeoryPool() {}

GLuint GLMemeoryPool::GLCreateSSBO(GLsizeiptr size, void *hostData, GLenum type, GLenum usage) {
    PanNN_ASSERT(size > 0);

    GLuint ssboBufferID;
    glGenBuffers(1, &ssboBufferID);
    OPENGL_CHECK_ERROR;

    glBindBuffer(type, ssboBufferID);
    OPENGL_CHECK_ERROR;
    PanNN_ASSERT(ssboBufferID > 0);

    glBufferData(type, size, hostData, usage);
    OPENGL_CHECK_ERROR;

    PanNN_ASSERT(m_ssboPool.count(ssboBufferID) == 0);
    m_ssboPool[ssboBufferID] = std::make_pair(size, type);

    return ssboBufferID;
}

ErrorCode GLMemeoryPool::CopyDeviceSSBOToHost(GLuint ssboBufferID, void *hostData, GLsizeiptr size) {
    PanNN_ASSERT(m_ssboPool.count(ssboBufferID) > 0);
    PanNN_ASSERT(m_ssboPool[ssboBufferID].first >= size);

    glBindBuffer(m_ssboPool[ssboBufferID].second, ssboBufferID);
    OPENGL_CHECK_ERROR;

    auto ptr = glMapBufferRange(m_ssboPool[ssboBufferID].second, 0, m_ssboPool[ssboBufferID].first, GL_MAP_READ_BIT);
    OPENGL_CHECK_ERROR;

    if (ptr != nullptr) {
        ::memcpy(hostData, ptr, size);
    }

    glUnmapBuffer(m_ssboPool[ssboBufferID].second);
    OPENGL_CHECK_ERROR;
    return NONE_ERR;
}

ErrorCode GLMemeoryPool::CopyHostToDeviceSSBO(void *hostData, GLuint ssboBufferID, GLsizeiptr size) {
    PanNN_ASSERT(m_ssboPool.count(ssboBufferID) > 0);
    PanNN_ASSERT(m_ssboPool[ssboBufferID].first >= size);

    glBindBuffer(m_ssboPool[ssboBufferID].second, ssboBufferID);
    OPENGL_CHECK_ERROR;

    auto ptr = glMapBufferRange(m_ssboPool[ssboBufferID].second, 0, m_ssboPool[ssboBufferID].first, GL_MAP_READ_BIT);
    OPENGL_CHECK_ERROR;

    if (ptr != nullptr) {
        ::memcpy(ptr, hostData, size);
    }

    glUnmapBuffer(m_ssboPool[ssboBufferID].second);
    OPENGL_CHECK_ERROR;
    return NONE_ERR;
}

GLuint GLMemeoryPool::GLCreateTexture(int w, int h, int c, GLenum textrueFormat, GLenum target) {
    GLuint textureID;
    if (target == GL_TEXTURE_3D) {
        PanNN_ASSERT(w > 0 && h > 0 && c > 0);
        glGenTextures(1, &textureID);
        OPENGL_CHECK_ERROR;
        glBindTexture(target, textureID);
        OPENGL_CHECK_ERROR;
        glTexParameteri(target, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
        OPENGL_CHECK_ERROR;
        glTexParameteri(target, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
        OPENGL_CHECK_ERROR;
        glTexParameteri(target, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
        OPENGL_CHECK_ERROR;
        glTexParameteri(target, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
        OPENGL_CHECK_ERROR;
        glTexParameteri(target, GL_TEXTURE_WRAP_R, GL_CLAMP_TO_EDGE);
        OPENGL_CHECK_ERROR;

        int realW = w;
        int realH = h;
        int realD = UP_DIV(c, 4);
        glTexStorage3D(target, 1, textrueFormat, realW, realH, realD);
        OPENGL_CHECK_ERROR;
    } else if (target == GL_TEXTURE_2D) {
        PanNN_ASSERT(w > 0 && h > 0);
        glGenTextures(1, &textureID);
        OPENGL_CHECK_ERROR;
        glBindTexture(target, textureID);
        OPENGL_CHECK_ERROR;
        glTexParameteri(target, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
        OPENGL_CHECK_ERROR;
        glTexParameteri(target, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
        OPENGL_CHECK_ERROR;
        glTexParameteri(target, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
        OPENGL_CHECK_ERROR;
        glTexParameteri(target, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
        OPENGL_CHECK_ERROR;
        glTexParameteri(target, GL_TEXTURE_WRAP_R, GL_CLAMP_TO_EDGE);
        OPENGL_CHECK_ERROR;

        int realW = w;
        int realH = h;
        glTexStorage2D(target, 1, textrueFormat, realW, realH);
        OPENGL_CHECK_ERROR;
    }

    std::vector<int> dims = {w, h, c};
    std::vector<GLenum> props = {textrueFormat, target};
    // m_texturePool[textureID] = std::make_pair({w, h, c}, {textrueFormat, target});
    m_texturePool[textureID] = std::make_pair(dims, props);
    return textureID;
}

ErrorCode GLMemeoryPool::CopyDeviceTextureToSSBO(GLuint textureID, GLuint ssboBufferID) {
    GLuint computeProgram = OpenGLRuntime::CreateComputeProgram(g_glsl_device_to_host);
    glUseProgram(computeProgram);

    // bind the src image texture
    glBindImageTexture(0, textureID, 0, GL_TRUE, 0, GL_READ_ONLY, GL_RGBA32F);

    // bind the dest output data
    glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, ssboBufferID);

    // set uniform values
    int width = m_texturePool[textureID].first[0];
    int height = m_texturePool[textureID].first[1];
    int channel = m_texturePool[textureID].first[2];

    glUniform1i(2, width);
    glUniform1i(3, height);
    glUniform1i(4, channel);

    int c_4 = UP_DIV(channel, 4);
    int gLocalSize[3] = {4, 4, 1};
    glDispatchCompute(UP_DIV(width, gLocalSize[0]), UP_DIV(height, gLocalSize[1]), UP_DIV(c_4, gLocalSize[2]));

    // memory sync
    glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);
    glDeleteProgram(computeProgram);
    return NONE_ERR;
}

ErrorCode GLMemeoryPool::CopyDeviceSSBOToTexture(GLuint ssboBufferID, GLuint textureID) {
    GLuint computeProgram = OpenGLRuntime::CreateComputeProgram(g_glsl_host_to_device);
    glUseProgram(computeProgram);

    // bind the src image texture
    glBindImageTexture(0, textureID, 0, GL_TRUE, 0, GL_WRITE_ONLY, GL_RGBA32F);

    // bind the dest output data
    glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, ssboBufferID);

    // set uniform values
    int width = m_texturePool[textureID].first[0];
    int height = m_texturePool[textureID].first[1];
    int channel = m_texturePool[textureID].first[2];

    glUniform1i(2, width);
    glUniform1i(3, height);
    glUniform1i(4, channel);

    int c_4 = UP_DIV(channel, 4);
    int gLocalSize[3] = {4, 4, 1};
    glDispatchCompute(UP_DIV(width, gLocalSize[0]), UP_DIV(height, gLocalSize[1]), UP_DIV(c_4, gLocalSize[2]));

    // memory sync
    glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);
    glDeleteProgram(computeProgram);
    return NONE_ERR;
}

GLuint GLMemeoryPool::CopyHostToDeviceTexture(void *hostData, int width, int height, int channel) {
    auto ssboBufferID = GLCreateSSBO(sizeof(float) * width * height * channel, hostData);
    auto textureID = GLCreateTexture(width, height, channel, GL_RGBA32F);
    CopyDeviceSSBOToTexture(ssboBufferID, textureID);
    return textureID;
}

void *GLMemeoryPool::CopyDeviceTextureToHost(GLuint textureID) {
    int width = m_texturePool[textureID].first[0];
    int height = m_texturePool[textureID].first[1];
    int channel = m_texturePool[textureID].first[2];

    auto ssboBufferID = GLCreateSSBO(sizeof(float) * width * height * channel);

    CopyDeviceTextureToSSBO(textureID, ssboBufferID);

    void *output = malloc(sizeof(float) * width * height * channel);
    CopyDeviceSSBOToHost(ssboBufferID, output, sizeof(float) * 256);
    return output;
}

}  // namespace OpenGL
}  // namespace PanNN
